import networkx as nx
import numpy as np
import random

import gen.feat as featgen
import util

def gen_ba(n_range, m_range, num_graphs, feature_generator=None):  #生成Barabási-Albert图（无标度网络）
    graphs = []
    for i in np.random.choice(n_range, num_graphs):
        for j in np.random.choice(m_range, 1):
            graphs.append(nx.barabasi_albert_graph(i,j))

    if feature_generator is None:
        feature_generator = ConstFeatureGen(0)
    for G in graphs:
        feature_generator.gen_node_features(G)
    return graphs

def gen_er(n_range, p, num_graphs, feature_generator=None):
    graphs = []
    for i in np.random.choice(n_range, num_graphs):
        graphs.append(nx.erdos_renyi_graph(i,p))

    if feature_generator is None:
        feature_generator = ConstFeatureGen(0)
    for G in graphs:
        feature_generator.gen_node_features(G)
    return graphs

def gen_2community_ba(n_range, m_range, num_graphs, inter_prob, feature_generators):
    ''' Each community is a BA graph.
    Args:
        inter_prob: probability of one node connecting to any node in the other community.
    '''

    if feature_generators is None:
        mu0 = np.zeros(10)
        mu1 = np.ones(10)
        sigma0 = np.ones(10, 10) * 0.1
        sigma1 = np.ones(10, 10) * 0.1
        fg0 = GaussianFeatureGen(mu0, sigma0)
        fg1 = GaussianFeatureGen(mu1, sigma1)
    else:
        fg0 = feature_generators[0]
        fg1 = feature_generators[1] if len(feature_generators) > 1 else feature_generators[0]

    graphs1 = []
    graphs2 = []
    #for (i1, i2) in zip(np.random.choice(n_range, num_graphs), 
    #                    np.random.choice(n_range, num_graphs)):
    #    for (j1, j2) in zip(np.random.choice(m_range, num_graphs), 
    #                        np.random.choice(m_range, num_graphs)):
    graphs0 = gen_ba(n_range, m_range, num_graphs, fg0)
    graphs1 = gen_ba(n_range, m_range, num_graphs, fg1)
    graphs = []
    for i in range(num_graphs):
        G = nx.disjoint_union(graphs0[i], graphs1[i])
        n0 = graphs0[i].number_of_nodes()
        for j in range(n0):
            if np.random.rand() < inter_prob:
                target = np.random.choice(G.number_of_nodes() - n0) + n0
                G.add_edge(j, target)
        graphs.append(G)
    return graphs

def gen_2hier(num_graphs, num_clusters, n, m_range, inter_prob1, inter_prob2, feat_gen):
    ''' Each community is a BA graph.
    Args:
        inter_prob1: probability of one node connecting to any node in the other community within
            the large cluster.
        inter_prob2: probability of one node connecting to any node in the other community between
            the large cluster.
    '''
    graphs = []

    for i in range(num_graphs):
        clusters2 = []
        for j in range(len(num_clusters)):
            clusters = gen_er(range(n, n+1), 0.5, num_clusters[j], feat_gen[0])
            G = nx.disjoint_union_all(clusters)
            for u1 in range(G.number_of_nodes()):
                if np.random.rand() < inter_prob1:
                    target = np.random.choice(G.number_of_nodes() - n)
                    # move one cluster after to make sure it's not an intra-cluster edge
                    if target // n >= u1 // n:
                        target += n
                    G.add_edge(u1, target)
            clusters2.append(G)
        G = nx.disjoint_union_all(clusters2)
        cluster_sizes_cum = np.cumsum([cluster2.number_of_nodes() for cluster2 in clusters2])
        curr_cluster = 0
        for u1 in range(G.number_of_nodes()):
            if u1 >= cluster_sizes_cum[curr_cluster]:
                curr_cluster += 1
            if np.random.rand() < inter_prob2:
                target = np.random.choice(G.number_of_nodes() -
                        clusters2[curr_cluster].number_of_nodes())
                # move one cluster after to make sure it's not an intra-cluster edge
                if curr_cluster == 0 or target >= cluster_sizes_cum[curr_cluster - 1]:
                    target += cluster_sizes_cum[curr_cluster]
            G.add_edge(u1, target)
        graphs.append(G)

    return graphs


def gen_ba_community(num_nodes=100, num_communities=5, community_size=None, 
                     m=3, inter_prob=0.05):
    """
    生成带社区结构的BA图（用于节点分类任务）
    
    每个社区是一个BA子图，社区间有少量连接
    
    参数:
        num_nodes (int): 总节点数
        num_communities (int): 社区数量
        community_size (int): 每个社区的大小（如果None，则平均分配）
        m (int): BA图参数
        inter_prob (float): 社区间连接概率
    
    返回:
        G: NetworkX图
        labels: 节点标签（社区ID）
    """
    if community_size is None:
        community_size = num_nodes // num_communities
    
    # 生成每个社区（BA子图）
    communities = []
    for i in range(num_communities):
        # 每个社区大小稍有变化
        size = community_size + np.random.randint(-5, 6)
        size = max(10, size)  # 至少10个节点
        
        # 生成BA子图
        G_sub = nx.barabasi_albert_graph(size, m)
        communities.append(G_sub)
    
    # 合并所有社区
    G = nx.disjoint_union_all(communities)
    
    # 生成节点标签（社区ID）
    labels = np.zeros(G.number_of_nodes(), dtype=np.int64)
    offset = 0
    for i, comm in enumerate(communities):
        size = comm.number_of_nodes()
        labels[offset:offset+size] = i
        offset += size
    
    # 添加社区间连接
    community_offsets = [0]
    for comm in communities:
        community_offsets.append(community_offsets[-1] + comm.number_of_nodes())
    
    for i in range(G.number_of_nodes()):
        # 找到节点i所在的社区
        comm_i = np.searchsorted(community_offsets[1:], i)
        
        # 以一定概率连接到其他社区的节点
        if np.random.rand() < inter_prob:
            # 随机选择另一个社区
            other_comm = np.random.choice([c for c in range(num_communities) if c != comm_i])
            # 随机选择该社区的一个节点
            start = community_offsets[other_comm]
            end = community_offsets[other_comm + 1]
            target = np.random.randint(start, end)
            G.add_edge(i, target)
    
    return G, labels


def gen_sbm_community(num_nodes=100, num_communities=5, p_in=0.3, p_out=0.01):
    """
    使用随机块模型（SBM）生成带社区结构的图
    
    参数:
        num_nodes (int): 总节点数
        num_communities (int): 社区数量
        p_in (float): 社区内连接概率
        p_out (float): 社区间连接概率
    
    返回:
        G: NetworkX图
        labels: 节点标签（社区ID）
    """
    community_size = num_nodes // num_communities
    sizes = [community_size] * num_communities
    
    # 调整最后一个社区以匹配总节点数
    sizes[-1] = num_nodes - sum(sizes[:-1])
    
    # 构造概率矩阵
    probs = np.ones((num_communities, num_communities)) * p_out
    np.fill_diagonal(probs, p_in)
    
    # 生成SBM图
    G = nx.stochastic_block_model(sizes, probs, seed=None)
    
    # 生成标签
    labels = np.zeros(num_nodes, dtype=np.int64)
    offset = 0
    for i, size in enumerate(sizes):
        labels[offset:offset+size] = i
        offset += size
    
    return G, labels

